#include <iostream>
#include <string>
#include <fstream>
#include <sstream>
#include <cstdlib>
#include <vector>
#include <sys/time.h>
#include <cuda.h>
#include <cuda_runtime_api.h>
#include <numeric>
#include <cassert>
using namespace std;

// nvcc -o coo_atom_new coo_atom_new.cu

void verify(int nz, int M, vector<unsigned int> &row_index, vector<unsigned int> &col_index, vector<float> &val, vector<float> &vec, vector<float> &out)
{
    float *correct = (float *)malloc(sizeof(float) * M);
    memset(correct, 0, sizeof(float) * M);
    for (int i = 0; i < nz; ++i)
    {
        correct[row_index[i]] += val[i] * vec[col_index[i]];
    }
    int count = 0;
    FILE *fp;
    fp = fopen("log.txt", "w+");
    for (int i = 0; i < M; i++)
    {

        if (fabs(correct[i] - out[i]) / correct[i] > 0.)
        {
            count++;
            fprintf(fp, "Output - %f, Correct - %f\n", out[i], correct[i]);
            fprintf(fp, "Row index is: %d\n", i);
        }
    }
    cout << "total wrong"
         << " " << count << endl;
}

void split(const string &s, vector<string> &sv, const char delim = ' ')
{
    sv.clear();
    istringstream iss(s);
    string temp;

    while (getline(iss, temp, delim))
    {
        sv.emplace_back(move(temp));
    }

    return;
}

void get_matrix_index_and_val_from_file(string coo_file_name, vector<unsigned int> &row_index_vec, vector<unsigned int> &col_index_vec, vector<float> &float_val_vec, unsigned int &max_row_index, unsigned int &max_col_index)
{
    max_col_index = 0;
    max_row_index = 0;

    // 读文件
    char buf[1024];

    ifstream infile;
    infile.open(coo_file_name);

    bool dataset_first_line = true;

    if (infile.is_open())
    {
        while (infile.good() && !infile.eof())
        {
            string line_str;
            vector<string> sv;
            memset(buf, 0, 1024);
            infile.getline(buf, 1024);
            line_str = buf;
            if (isspace(line_str[0]) || line_str.empty() || line_str[0] == '%')
            {
                continue;
            }

            split(line_str, sv);

            if (dataset_first_line == true)
            {
                max_row_index = atol(sv[0].c_str()) - 1;
                max_col_index = atol(sv[1].c_str()) - 1;
                dataset_first_line = false;
                continue;
            }

            unsigned long row_index = atol(sv[0].c_str()) - 1;
            unsigned long col_index = atol(sv[1].c_str()) - 1;

            row_index_vec.push_back(row_index);
            col_index_vec.push_back(col_index);
            float_val_vec.push_back(1);

            if (row_index > max_row_index)
            {
                max_row_index = row_index;
            }

            if (col_index > max_col_index)
            {
                max_col_index = col_index;
            }
        }
    }
    else
    {
        cout << "get_matrix_index_and_val_from_file: cannot open file " << coo_file_name << endl;
    }

    infile.close();
    cout << "finish read file" << endl;
}
__global__ void spmv_coo_kernel(float *d_val, float *d_vector, unsigned int *d_cols, unsigned int *d_rows, float *d_out, unsigned int thread_nnz, unsigned int needed_thread_num, unsigned int nnz)
{
    unsigned int real_thread_num = gridDim.x * blockDim.x;
    unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;
    for (unsigned int i = tid; i < needed_thread_num; i = i + real_thread_num)
    {
        float temp_result = 0;
        unsigned int id = i;
        for (unsigned int j = 0; j < thread_nnz; j++)
        {
            temp_result += __ldg(&d_vector[d_cols[id]]) * d_val[id];
            id += needed_thread_num;
        }
        atomicAdd(&(d_out[d_rows[i]]), temp_result);
    }
}

void conversion(vector<unsigned int> &row_index, vector<unsigned int> &col_index, vector<float> &val, vector<unsigned int> &new_row_index, vector<unsigned int> &new_col_index, vector<float> &new_val, int row_num, int thread_nnz)
{

    unsigned int *nnz_of_each_row;
    nnz_of_each_row = (unsigned int *)calloc(row_num, sizeof(unsigned int));

    for (unsigned long cur_nz_index = 0; cur_nz_index < row_index.size(); cur_nz_index++)
    {
        unsigned long cur_row_index = row_index[cur_nz_index];
        nnz_of_each_row[cur_row_index] = nnz_of_each_row[cur_row_index] + 1;
    }
    unsigned int total_size = 0;
    for (unsigned int i = 0; i < row_num; i++)
    {
        if (nnz_of_each_row[i] == 0)
        {
            continue;
        }
        total_size += (nnz_of_each_row[i] % thread_nnz == 0) ? nnz_of_each_row[i] : (nnz_of_each_row[i] / thread_nnz + 1) * thread_nnz;
    }
    vector<unsigned int> row(total_size);
    vector<unsigned int> col(total_size);
    vector<float> value(total_size, 0);
    
    unsigned int cur_total_id = 0;
    unsigned int cur_total_id_padding = 0;
    for (unsigned int j = 0; j < row_num; j++)
    {
        int cur_len = nnz_of_each_row[j];
        int target_len = cur_len % thread_nnz == 0 ? cur_len : (cur_len / thread_nnz + 1) * thread_nnz;
        if (cur_len != 0)
        {
            for (int id = 0; id < cur_len; id++)
            {
                row[cur_total_id_padding + id] = row_index[cur_total_id + id];
                col[cur_total_id_padding + id] = col_index[cur_total_id + id];
                value[cur_total_id_padding + id] = val[cur_total_id + id];
            }
            int row_last = row_index[cur_total_id + cur_len - 1];
            int col_last = col_index[cur_total_id + cur_len - 1];
            for (int add_id = 0; add_id < target_len - cur_len; add_id++)
            {
                row[cur_total_id_padding + cur_len + add_id] = row_last;
                col[cur_total_id_padding + cur_len + add_id] = col_last;
            }

            cur_total_id_padding += target_len;
            cur_total_id += cur_len;
        }
    }

    vector<unsigned int> interlance_row(total_size);
    vector<unsigned int> interlance_col(total_size);
    vector<float> interlance_val(total_size);
    unsigned int thread_num = total_size / thread_nnz;

    for (unsigned int thread_id = 0; thread_id < thread_num; thread_id++)
    {
        for (int id_in_thread = 0; id_in_thread < thread_nnz; id_in_thread++)
        {
            interlance_row[thread_id + id_in_thread * thread_num] = row[thread_id * thread_nnz + id_in_thread];
            interlance_col[thread_id + id_in_thread * thread_num] = col[thread_id * thread_nnz + id_in_thread];
            interlance_val[thread_id + id_in_thread * thread_num] = value[thread_id * thread_nnz + id_in_thread];
        }
    }
    new_row_index = interlance_row;
    new_col_index = interlance_col;
    new_val = interlance_val;
}

void spmv_coo(vector<unsigned int> &row_index, vector<unsigned int> &col_index, vector<float> &val, unsigned int row_num, vector<float> &vec, vector<float> &out, int thread_nnz, int nblock, int block_size)
{
    vector<unsigned int> new_row_index, new_col_index;
    vector<float> new_val;
    struct timeval start, end;

    gettimeofday(&start, NULL);
    
    conversion(row_index, col_index, val, new_row_index, new_col_index, new_val, row_num, thread_nnz);
    
    gettimeofday(&end, NULL);

    long timeuse_ = 1000000 * (end.tv_sec - start.tv_sec) + end.tv_usec - start.tv_usec;
    printf("conversion time=%fms \n", timeuse_ / 1000.0);

    float *d_vec, *d_val, *d_out;
    unsigned int *d_rows, *d_cols;
    unsigned int nnz = new_val.size();
    unsigned int needed_thread_num = nnz / thread_nnz;
    int device;
    cudaGetDevice(&device);
    cudaMalloc(&d_vec, vec.size() * sizeof(float));
    cudaMalloc(&d_val, new_val.size() * sizeof(float));
    cudaMalloc(&d_out, out.size() * sizeof(float));
    cudaMalloc(&d_cols, new_col_index.size() * sizeof(unsigned int));
    cudaMalloc(&d_rows, new_row_index.size() * sizeof(unsigned int));

    cudaMemcpy(d_vec, &vec[0], vec.size() * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_val, &new_val[0], new_val.size() * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_cols, &new_col_index[0], new_col_index.size() * sizeof(unsigned int), cudaMemcpyHostToDevice);
    cudaMemcpy(d_rows, &new_row_index[0], new_row_index.size() * sizeof(unsigned int), cudaMemcpyHostToDevice);

    unsigned int repeat_num = 5;
    // struct timeval start, end;
    gettimeofday(&start, NULL);
    for (int repeat_id = 0; repeat_id < repeat_num; repeat_id++)
    {
        cudaMemset(d_out, 0, out.size() * sizeof(float));
        spmv_coo_kernel<<<nblock, block_size>>>(d_val, d_vec, d_cols, d_rows, d_out, thread_nnz, needed_thread_num, nnz);
        cudaDeviceSynchronize();
    }
    gettimeofday(&end, NULL);

    long timeuse = 1000000 * (end.tv_sec - start.tv_sec) + end.tv_usec - start.tv_usec;
    double gflops = ((double)2.0 * val.size() * repeat_num / ((double)timeuse / 1000000)) / 1000000000;

    printf("time=%fms, gflops=%f\n", timeuse / 1000.0, gflops);

    cudaMemcpy(&out[0], d_out, out.size() * sizeof(float), cudaMemcpyDeviceToHost);

    cudaFree(d_vec);
    cudaFree(d_val);
    cudaFree(d_cols);
    cudaFree(d_rows);
    cudaFree(d_out);
}

int main(int argc, char **argv)
{
    string file_name = argv[1];
    cout << file_name << endl;
    int block_num = std::atoi(argv[2]);
    int thread_num = std::atoi(argv[3]);
    int element_num = std::atoi(argv[4]);

    struct timeval start, end;
    // gettimeofday(&start, NULL);

    vector<unsigned int> row_index;
    vector<unsigned int> col_index;
    // vector<float> val;

    unsigned int row_num = 0;
    unsigned int col_num = 0;
    unsigned int nnz_num = 0;


    ////////////////////////

    // vector<unsigned int> row_index_;
    // vector<unsigned int> col_index_;
    // vector<float> val_;
    // unsigned int max_col_index_ = 0;
    // unsigned int max_row_index_ = 0;
    // string txt_file_name = file_name + ".mtx";
    
    // get_matrix_index_and_val_from_file(txt_file_name, row_index_, col_index_, val_, max_row_index_, max_col_index_);
    // vector<unsigned int> ori_row(row_index_);
    // vector<unsigned int> ori_col(col_index_);
    // vector<float> ori_val(val_);
    
    /////////////////

    string out_name = file_name + "_binary.mtx";
    ifstream inFile(out_name, ios::in|ios::binary); 
    
    inFile.read((char *)&row_num, sizeof(unsigned int));
    inFile.read((char *)&col_num, sizeof(unsigned int));
    inFile.read((char *)&nnz_num, sizeof(unsigned int));

    row_index.resize(nnz_num);
    col_index.resize(nnz_num);
    // val.resize(nnz_num);

    inFile.read((char *)&row_index[0], nnz_num * sizeof(unsigned int));
    inFile.read((char *)&col_index[0], nnz_num * sizeof(unsigned int));
    // inFile.read((char *)&val[0], nnz_num * sizeof(float));
    
    vector<float> val(nnz_num, 1.0);

    cout << "finish read file" << endl;
    
    // printf("%u %u\n", max_col_index, max_row_index);
    // printf("%u %u\n", max_col_index_, max_row_index_);
    // printf("%u nnz_size %lu\n", nnz_num, val.size());
    // for (int i = 0; i < nnz_num; i++)
    // {
    //     if (row_index[i] != row_index_[i])
    //     {
    //         printf("row: %d\n", i);
    //         std::getchar();
    //     }
    //     if (col_index[i] != col_index_[i])
    //     {
    //         printf("col: %d\n", i);
    //         std::getchar();
    //     }
    //     if (val[i] != val_[i])
    //     {
    //         printf("value: %d\n", i);
    //         printf("value: %f %f\n", val[i], val_[i]);
    //         std::getchar();
    //     }
    // }

    vector<float> vec;
    vector<float> out;
    for (int row_id = 0; row_id < row_num; row_id++)
    {
        out.push_back(0);
    }

    for (int col_id = 0; col_id < col_num; col_id++)
    {
        vec.push_back(1);
    }

    // gettimeofday(&end, NULL);

    // long timeuse = 1000000 * (end.tv_sec - start.tv_sec) + end.tv_usec - start.tv_usec;
    // printf("data read: time=%fms \n", timeuse / 1000.0);

    gettimeofday(&start, NULL);

    printf("block_num: %d  thread_num: %d element_num: %d \n", block_num, thread_num, element_num);
    spmv_coo(row_index, col_index, val, row_num, vec, out, element_num, block_num, thread_num);

    gettimeofday(&end, NULL);
    long timeuse = 1000000 * (end.tv_sec - start.tv_sec) + end.tv_usec - start.tv_usec;
    printf("profiling whole time=%fms \n", timeuse / 1000.0);

    // verify(ori_row.size(), row_num, ori_row, ori_col, ori_val, vec, out);
}
